import json
import openai
import asyncio
import aiohttp
from datetime import datetime
from collections import defaultdict
from .models import Agent, Workflow, WorkflowExecution

class WorkflowExecutor:
    def __init__(self, workflow_id):
        self.workflow = Workflow.objects.get(id=workflow_id)
        self.node_map = {node['id']: node for node in self.workflow.nodes}
        self.edge_map = defaultdict(list)
        self.reverse_edge_map = defaultdict(list)
        for edge in self.workflow.edges:
            self.edge_map[edge['source']].append(edge['target'])
            self.reverse_edge_map[edge['target']].append(edge)
    
    def validate_workflow(self):
        """验证工作流结构有效性"""
        input_nodes = [n for n in self.workflow.nodes if n['type'] == 'inputNode']
        if len(input_nodes) == 0:
            raise ValueError("工作流必须包含输入节点")
            
        output_nodes = [n for n in self.workflow.nodes if n['type'] == 'outputNode']
        if len(output_nodes) == 0:
            raise ValueError("工作流必须包含输出节点")

    async def execute(self, input_data):
        """执行工作流"""
        context = {}
        
        # 找到起始输入节点
        input_nodes = [n for n in self.workflow.nodes if n['type'] == 'inputNode']
        
        # 异步执行所有分支
        tasks = [self.process_node(node['id'], input_data, context) 
                for node in input_nodes]
        await asyncio.gather(*tasks)
        
        # 收集输出节点结果
        output_nodes = [n for n in self.workflow.nodes if n['type'] == 'outputNode']
        return {
            node['id']: context.get(node['id'])
            for node in output_nodes
        }

    def execute_sync(self, input_data):
        """同步执行方法"""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            result = loop.run_until_complete(self.execute(input_data))
            return result
        finally:
            loop.close()

    async def process_node(self, node_id, input_data, context):
        """处理单个节点"""
        node = self.node_map[node_id]
        print('处理单个节点：', node)
        try:
            # 执行节点逻辑
            if node['type'] == 'inputNode':
                result = self.handle_input_node(node, input_data)
            elif node['type'] == 'agentNode':
                result = await self.handle_agent_node(node, context)
            elif node['type'] == 'outputNode':
                result = self.handle_output_node(node, context)
            elif node['type'] == 'httpNode':
                result = await self.handle_http_node(node, context)
            
            context[node_id] = result
            
            # 触发下游节点
            for target in self.edge_map.get(node_id, []):
                await self.process_node(target, input_data, context)
                
        except Exception as e:
            context[node_id] = {'error': str(e)}
            raise
    def exec_string(self, code, function_name, *args, **kwargs):
        # 创建一个空字典作为命名空间，并包含__builtins__
        namespace = {'__builtins__': __builtins__}
        try:
            # 执行字符串中的代码
            exec(code, namespace)
            
            # 从命名空间中获取函数并执行
            if function_name in namespace:
                func = namespace[function_name]
                if len(args) == 1 and isinstance(args[0], dict) and not args[0] and not kwargs:
                    result = func()
                else:
                    result = func(*args, **kwargs)
                return result
            else:
                return f'函数{function_name}未定义'
        except Exception as e:
            return f'执行出错: {e}'

    async def call_http_tool(self, config, data):
        try:
            # 提取 HTTP 配置
            url = config['httpConfig']['url']
            method = config['httpConfig']['method']
            body_data = {}
            headers = {}
            for head in config['httpConfig']['headers']:
                headers[head['hkey']] = head['hvalue']
            if not url:
                raise ValueError("URL is required in httpConfig")

            async with aiohttp.ClientSession() as session:
                # 动态调用不同的 HTTP 方法
                async with session.request(method, url, json=body_data, headers=headers, params=data) as response:
                    # 处理返回值
                    if response.status == 200:
                        # 返回 JSON 数据或文本内容
                        if response.content_type == 'application/json':
                            result = await response.json()
                        else:
                            result = await response.text()
                    else:
                        result = {
                            "status": response.status,
                            "error": await response.text()
                        }
            print("HTTP call result:", result)
            return result

        except Exception as e:
            print(f"Error during HTTP call: {e}")
            return f"Error during HTTP call: {e}"
   
    def handle_input_node(self, node, input_data):
        """处理输入节点"""
        # 根据节点配置映射输入数据
        print('input_data: ', input_data)
        print(node)
        return input_data.get(node['id'], '')

    async def handle_agent_node(self, node, context):
        """处理LLM节点"""
        config = node['data']
        print('llm node config: ', node)
        
        # 获取上游数据
        inputs = []
        for edge in self.reverse_edge_map.get(node['id'], []):
            if edge.get('targetHandle') != 'tools':
                upstream_data = context.get(edge['source'], {})
                inputs.append(upstream_data)

        # 收集工具节点并转换为函数工具定义
        tool_definitions = []
        functions = {}
        for edge in self.reverse_edge_map.get(node['id'], []):
            if edge.get('targetHandle') == 'tools' and edge.get('sourceHandle') == 'tool':
                tool_node = self.node_map.get(edge['source'])
                if tool_node and tool_node['type'] == 'toolNode':
                    tool_data = tool_node['data']
                    functions[tool_data['toolName']] = tool_data

                    parameters = {
                        'type': 'object',
                        'properties': {},
                        'required': []
                    }
                    for param in tool_data.get('params', []):
                        param_name = param['name']
                        parameters['properties'][param_name] = {
                            'type': param.get('type', 'string'),
                            'description': param.get('description', '')
                        }
                        parameters['required'].append(param_name)
                    # 去重required参数
                    parameters['required'] = list(set(parameters['required']))
                    
                    function_def = {
                        'name': tool_data['toolName'],
                        'description': tool_data.get('desc', ''),
                        'parameters': parameters
                    }
                    tool_definitions.append({
                        'type': 'function',
                        'function': function_def
                    })
        print('tool_definitions: \n', tool_definitions)
        print('functions: \n', functions)
        print('config: \n', config)
        
        # 调用LLM服务（示例实现）
        # url = config.get('apiUrl')
        # appkey = config.get('appKey')
        return await self.call_llm_service(
            inputs=inputs,
            config = config,
            tools=tool_definitions,
            functions=functions
        )

    async def call_llm_service(self, inputs, config, tools=None, functions=None):
        """调用实际LLM服务"""
        client = openai.OpenAI(
                api_key=config['apiKey'],
                base_url=config['apiUrl'],
            )
        messages=[
            # {'role': 'system', 'content': config['desc']},
            {"role": "user", "content": '\n'.join(inputs)}
        ]
        if tools:
            response = client.chat.completions.create(
                model=config['modelName'],
                messages=messages,
                temperature=float(config['temperature']),
                max_tokens=int(config['maxPrompt']),
                tools=tools,
                tool_choice='auto'
            )
        else:
            response = client.chat.completions.create(
                    model=config['modelName'],
                    messages=messages,
                    temperature=float(config['temperature']),
                    max_tokens=int(config['maxPrompt'])
                )
        print('response.choices[0].message: \n', response.choices[0].message)
        tool_calls = response.choices[0].message.tool_calls or []
        
        tool_results = []
        for tool_call in tool_calls:
            tool_name = tool_call.function.name
            parameters = json.loads(tool_call.function.arguments)
            print('工具名称与参数：', tool_name, parameters)
            if functions[tool_name]['execType'] == 'Python 代码':
                tool_result = self.exec_string(functions[tool_name]['codeContent'], tool_name, parameters)
            else:
                tool_result = await self.call_http_tool(functions[tool_name], parameters)
            print('tool_result: ', tool_result)
            tool_results.append({
                'tool_name': tool_name,
                'arguments': parameters,
                'content': tool_result
            })
        tool_message = {
                'role': 'assistant',
                'content': None,
                'tool_calls': [
                    {
                        'id': f'result_{index}',
                        'type': 'function',
                        'function': {
                            'name': res['tool_name'],
                            'arguments': json.dumps(res['arguments'])
                        }
                    } for index, res in enumerate(tool_results)
                ]
            }

        if len(tool_results) > 0:
            messages.append(tool_message)
            for index, res in enumerate(tool_results):
                messages.append({
                    'role': 'tool',
                    'content': json.dumps(res['content']),
                    'tool_call_id': f'result_{index}'
                })
            print('messages tool: \n', messages)
            response = client.chat.completions.create(
                    model=config['modelName'],
                    messages=messages,
                    temperature=float(config['temperature']),
                    max_tokens=int(config['maxPrompt'])
                )
        
        print('response.choices[0].message.content: \n', response.choices[0].message.content)
        return {
            'content': response.choices[0].message.content,
            # 'tool_calls': response.choices[0].message.tool_calls
        }

        print('response.choices[0].message: \n', response.choices[0].message)
        return response.choices[0].message
        # 示例：模拟异步API调用
        await asyncio.sleep(1)  
        return {
            'result': f"Processed by {model}: {prompt} | Inputs: {inputs}"
        }

    def handle_output_node(self, node, context):
        """处理输出节点"""
        # 收集所有上游数据
        outputs = []
        for edge in self.workflow.edges:
            if edge['target'] == node['id']:
                outputs.append(context.get(edge['source'], {}))
        print('outputs: \n', {
            'outputs': outputs,
            'timestamp': datetime.now().isoformat()
        })    
        return {
            'outputs': outputs,
            'timestamp': datetime.now().isoformat()
        }

    async def handle_http_node(self, node, context):
        """处理http节点"""
        config = node['data']
        # 获取上游数据
        inputs = []
        for edge in self.reverse_edge_map.get(node['id'], []):
                upstream_data = context.get(edge['source'], {})
                inputs.append(upstream_data)
        data = inputs[0]
        try:
            # 提取 HTTP 配置
            url = config['url']
            method = config['method']
            body_data = {}
            headers = {}
            for head in config['headers']:
                headers[head['key']] = head['value']
            if not url:
                raise ValueError("URL is required in httpConfig")

            async with aiohttp.ClientSession() as session:
                # 动态调用不同的 HTTP 方法
                async with session.request(method, url, json=body_data, headers=headers, params=data) as response:
                    # 处理返回值
                    if response.status == 200:
                        # 返回 JSON 数据或文本内容
                        if response.content_type == 'application/json':
                            result = await response.json()
                        else:
                            result = await response.text()
                    else:
                        result = {
                            "status": response.status,
                            "error": await response.text()
                        }
            print("HTTP call result:", result)
            return result

        except Exception as e:
            print(f"Error during HTTP call: {e}")
            return f"Error during HTTP call: {e}"
